import os
import argparse
import torch
import cv2
import numpy as np
import time
from libs.utils import getimages, hairrecolor
import wget
import importlib

#modify here
#-------------------------------------#
def preprocess(img):
    img = img.astype(np.float32)
    img = img / 255.
    img = (img-0.5)/0.5
    img = np.transpose(img, (2,0,1))
    img = np.expand_dims(img, axis=0)
    return img
#-------------------------------------#


def evaluate(net_name, model_path, img_folder, input_shape):
    module  =importlib.import_module("nets.%s"%net_name)
    model = module.get_segmatting_model()
    model.load_state_dict(torch.load(model_path))
    model = model.to("cpu")
    model.eval()
    for bgr_img, rgb_img in getimages(img_folder):
        rgb_img = cv2.resize(rgb_img, input_shape)
        processimg = preprocess(rgb_img)
        input_data = torch.from_numpy(processimg).to("cpu")
        input_data = input_data.float()

        #modify here
        #----------------------------------------#
        mask, alpha = model(input_data)
        mask = mask.detach().numpy()[0]
        h, w = bgr_img.shape[0:2]
        mask = mask.argmax(axis=0)
        mask = mask.astype(np.uint8)
        alpha =  alpha.detach().numpy()[0][0]
        dyeimg = hairrecolor(bgr_img, mask, (0x40, 0x16, 0x66))
        cv2.imshow("img", cv2.resize(dyeimg, (500, 500)))
        cv2.waitKey(0)
        #----------------------------------------#

def main():

    parser = argparse.ArgumentParser()
    parser.add_argument("--net_name", default='ESPNet')
    parser.add_argument("--model_path", default='')
    parser.add_argument("--oss_model_url", default='https://dlmodels.oss-cn-beijing.aliyuncs.com/hairMatting_pytorch/hairnet.pth')
    parser.add_argument("--oss_model_path", default='./sample/models/hairnet.pth')
    parser.add_argument('--img_folder', default='./sample/pictures')
    parser.add_argument('--input_shape', default=(480, 480) )
    args = parser.parse_args()

    if args.model_path is '':
        if not os.path.exists(args.oss_model_path):
            wget.download(args.oss_model_url, out=args.oss_model_path)
        model_path = args.oss_model_path
    else:
        model_path = args.model_path

    evaluate(args.net_name, model_path, args.img_folder, args.input_shape)


if __name__ == "__main__":
    main()
